import csv
import argparse
from collections import defaultdict

# 读取BED文件，返回包含(chromosome, pos, strand)的集合
def read_bed_file(file_path):
    bed_set = set()
    with open(file_path, 'r') as bed_file:
        reader = csv.reader(bed_file, delimiter='\t')
        for row in reader:
            if len(row) >= 3:
                bed_set.add((row[0], int(row[1]), row[2]))
    return bed_set

# 读取TSV文件，并根据正样本和负样本分别统计
def analyze_tsv_file(tsv_path, positive_bed, negative_bed):
    positive_stats = {
        'prob_1_le_0.5': 0,
        'prob_1_le_0.5_seqs': defaultdict(int),
        'prob_1_0.5_0.6': 0,
        'prob_1_0.5_0.6_seqs': defaultdict(int)
    }
    
    negative_stats = {
        'prob_0_le_0.5': 0,
        'prob_0_le_0.5_seqs': defaultdict(int),
        'prob_0_0.5_0.6': 0,
        'prob_0_0.5_0.6_seqs': defaultdict(int)
    }
    
    with open(tsv_path, 'r') as tsv_file:
        reader = csv.reader(tsv_file, delimiter='\t')
        next(reader)  # 跳过标题行，如果有的话
        for row in reader:
            if len(row) >= 10:
                chromosome, pos, strand = row[0], int(row[1]), row[2]
                prob_0, prob_1 = float(row[6]), float(row[7])
                seq = row[9]
                key = (chromosome, pos, strand)
                
                # 对正样本进行统计
                if key in positive_bed:
                    if prob_1 <= 0.5:
                        positive_stats['prob_1_le_0.5'] += 1
                        positive_stats['prob_1_le_0.5_seqs'][seq] += 1
                    elif 0.5 < prob_1 <= 0.6:
                        positive_stats['prob_1_0.5_0.6'] += 1
                        positive_stats['prob_1_0.5_0.6_seqs'][seq] += 1
                
                # 对负样本进行统计
                if key in negative_bed:
                    if prob_0 <= 0.5:
                        negative_stats['prob_0_le_0.5'] += 1
                        negative_stats['prob_0_le_0.5_seqs'][seq] += 1
                    elif 0.5 < prob_0 <= 0.6:
                        negative_stats['prob_0_0.5_0.6'] += 1
                        negative_stats['prob_0_0.5_0.6_seqs'][seq] += 1
    
    return positive_stats, negative_stats

# 打印统计结果
def print_stats(stats, label):
    print(f"\n{label} Statistics:")
    # 打印满足prob <= 0.5的记录和序列信息
    prob_le_0_5_key = 'prob_1_le_0.5' if 'prob_1_le_0.5' in stats else 'prob_0_le_0.5'
    prob_le_0_5_seqs_key = 'prob_1_le_0.5_seqs' if 'prob_1_le_0.5_seqs' in stats else 'prob_0_le_0.5_seqs'
    
    print(f"Number of records with prob <= 0.5: {stats[prob_le_0_5_key]}")
    print(f"Number of unique sequences in prob <= 0.5: {len(stats[prob_le_0_5_seqs_key])}")
    print("Sequences and their counts in prob <= 0.5:")
    for seq, count in stats[prob_le_0_5_seqs_key].items():
        print(f"  {seq}: {count}")
    
    # 打印满足0.5 < prob <= 0.6的记录和序列信息
    prob_0_5_0_6_key = 'prob_1_0.5_0.6' if 'prob_1_0.5_0.6' in stats else 'prob_0_0.5_0.6'
    prob_0_5_0_6_seqs_key = 'prob_1_0.5_0.6_seqs' if 'prob_1_0.5_0.6_seqs' in stats else 'prob_0_0.5_0.6_seqs'
    
    print(f"\nNumber of records with 0.5 < prob <= 0.6: {stats[prob_0_5_0_6_key]}")
    print(f"Number of unique sequences in 0.5 < prob <= 0.6: {len(stats[prob_0_5_0_6_seqs_key])}")
    print("Sequences and their counts in 0.5 < prob <= 0.6:")
    for seq, count in stats[prob_0_5_0_6_seqs_key].items():
        print(f"  {seq}: {count}")

def main():
    # 设置命令行参数解析
    parser = argparse.ArgumentParser(description="统计TSV文件中与正负样本BED匹配的记录")
    parser.add_argument("positive_bed_path", help="正样本BED文件路径")
    parser.add_argument("negative_bed_path", help="负样本BED文件路径")
    parser.add_argument("tsv_path", help="TSV文件路径")
    args = parser.parse_args()

    # 读取正样本和负样本BED文件
    positive_bed = read_bed_file(args.positive_bed_path)
    negative_bed = read_bed_file(args.negative_bed_path)

    # 读取TSV文件并进行分析
    positive_stats, negative_stats = analyze_tsv_file(args.tsv_path, positive_bed, negative_bed)

    # 打印正样本和负样本的统计结果
    print_stats(positive_stats, "Positive Sample")
    print_stats(negative_stats, "Negative Sample")

if __name__ == "__main__":
    main()
